import org.bouncycastle.pqc.math.linearalgebra.ByteUtils;

import javax.crypto.BadPaddingException;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.util.Arrays;

public class sm4_demo extends GMBaseTest {

    public static void main(String[] args) throws Exception {
        testEncryptAndDecrypt();
        testMac();
    }

    public static void testEncryptAndDecrypt() throws
            NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException,
            IllegalBlockSizeException, BadPaddingException, InvalidKeyException,
            InvalidAlgorithmParameterException {
        byte[] key = SM4Util.generateKey();
        byte[] iv = SM4Util.generateKey();
        byte[] cipherText = null;
        byte[] decryptedData = null;

        cipherText = SM4Util.encrypt_ECB_NoPadding(key, SRC_DATA_16B);
        System.out.println("SM4 ECB NoPadding encrypt result:\n" + Arrays.toString(cipherText));
        decryptedData = SM4Util.decrypt_ECB_NoPadding(key, cipherText);
        System.out.println("SM4 ECB NoPadding decrypt result:\n" + Arrays.toString(decryptedData));
        if (!Arrays.equals(decryptedData, SRC_DATA_16B)) {
            System.out.println("解密结果与明文不一致。\n");
        }

        cipherText = SM4Util.encrypt_ECB_Padding(key, SRC_DATA);
        System.out.println("SM4 ECB Padding encrypt result:\n" + Arrays.toString(cipherText));
        decryptedData = SM4Util.decrypt_ECB_Padding(key, cipherText);
        System.out.println("SM4 ECB Padding decrypt result:\n" + Arrays.toString(decryptedData));
        if (!Arrays.equals(decryptedData, SRC_DATA)) {
            System.out.println("解密结果与明文不一致。\n");
        }

        cipherText = SM4Util.encrypt_CBC_Padding(key, iv, SRC_DATA);
        System.out.println("SM4 CBC Padding encrypt result:\n" + Arrays.toString(cipherText));
        decryptedData = SM4Util.decrypt_CBC_Padding(key, iv, cipherText);
        System.out.println("SM4 CBC Padding decrypt result:\n" + Arrays.toString(decryptedData));
        if (!Arrays.equals(decryptedData, SRC_DATA)) {
            System.out.println("解密结果与明文不一致。\n");
        }

        cipherText = SM4Util.encrypt_CBC_NoPadding(key, iv, SRC_DATA_16B);
        System.out.println("SM4 CBC NoPadding encrypt result:\n" + Arrays.toString(cipherText));
        decryptedData = SM4Util.decrypt_CBC_NoPadding(key, iv, cipherText);
        System.out.println("SM4 CBC NoPadding decrypt result:\n" + Arrays.toString(decryptedData));
        if (!Arrays.equals(decryptedData, SRC_DATA_16B)) {
            System.out.println("解密结果与明文不一致。\n");
        }

    }

    public static void testMac() throws Exception {
        byte[] key = SM4Util.generateKey();
        byte[] iv = SM4Util.generateKey();

        byte[] mac = SM4Util.doCMac(key, SRC_DATA_24B);
        System.out.println("CMAC:\n" + ByteUtils.toHexString(mac).toUpperCase());

        mac = SM4Util.doGMac(key, iv, 16, SRC_DATA_24B);
        System.out.println("GMAC:\n" + ByteUtils.toHexString(mac).toUpperCase());

        byte[] cipher = SM4Util.encrypt_CBC_NoPadding(key, iv, SRC_DATA_32B);
        byte[] cipherLast16 = Arrays.copyOfRange(cipher, cipher.length - 16, cipher.length);
        mac = SM4Util.doCBCMac(key, iv, null, SRC_DATA_32B);
        if (!Arrays.equals(cipherLast16, mac)) {
            System.out.println("MAC验证不一致。\n");
        }
        System.out.println("CBCMAC:\n" + ByteUtils.toHexString(mac).toUpperCase());

        cipher = SM4Util.encrypt_CBC_Padding(key, iv, SRC_DATA_32B);
        cipherLast16 = Arrays.copyOfRange(cipher, cipher.length - 16, cipher.length);
        mac = SM4Util.doCBCMac(key, iv, SRC_DATA_32B);
        if (!Arrays.equals(cipherLast16, mac)) {
            System.out.println("MAC验证不一致。\n");
        }
        System.out.println("CBCMAC:\n" + ByteUtils.toHexString(mac).toUpperCase());
    }
}